package com.we.risk.sa;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.FilterFunction;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.connector.kafka.sink.KafkaRecordSerializationSchema;
import org.apache.flink.connector.kafka.sink.KafkaSink;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.contrib.streaming.state.RocksDBStateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.flink.util.Collector;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.we.flink.utils.WeKafkaPropertyReader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URLDecoder;
import java.util.Base64;
import java.util.Iterator;
import java.util.Properties;
import java.util.zip.GZIPInputStream;

public class RawSaEventAccessLogParseFilter {
    public static Logger LOG = LoggerFactory.getLogger(RawSaEventAccessLogParseFilter.class);
    public static final String GZIP_ENCODE_UTF_8 = "UTF-8";
    public static final String KEYBY = "distinct_id";
    public static final String RELEASEPROP = "risk/saevent/rawsaevent_kfk_filter.properties";

    public static void main(String[] args) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        try {
            WeKafkaPropertyReader paramReader = WeKafkaPropertyReader.init(RELEASEPROP);
            /** RocksDB configure */
            LOG.warn("no rocksdb url: " + paramReader.getRocksDBBackendUrl());
            env.setStateBackend(new RocksDBStateBackend(paramReader.getRocksDBBackendUrl()));

            /** checkpoint configure */
            CheckpointConfig ckConf = env.getCheckpointConfig();
            ckConf.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            ckConf.setCheckpointInterval(10 * 60 * 1000); // ms
            ckConf.setCheckpointTimeout(60 * 60 * 1000);
            ckConf.setMaxConcurrentCheckpoints(1);
            ckConf.setMinPauseBetweenCheckpoints(500);
            ckConf.enableExternalizedCheckpoints(
                    CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);

            ckConf.enableUnalignedCheckpoints();

            /** Consumer Kafka */
            String srcKafkaTopic = paramReader.getTmpKfkTopic();
            String srcKafkaBootStrapServer = paramReader.getTmpKfkBootStrapServer();
            String srckfkGrupId = paramReader.getTmpKfkGroupId();

            Properties consumProp = new Properties();
            consumProp.setProperty("bootstrap.servers", srcKafkaBootStrapServer);
            consumProp.setProperty("group.id", srckfkGrupId);
            consumProp.setProperty(
                    "key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
            consumProp.setProperty(
                    "value.deserializer",
                    "org.apache.kafka.common.serialization.StringDeserializer");

            KafkaSource<String> kfkSource =
                    KafkaSource.<String>builder()
                            .setTopics(srcKafkaTopic)
                            .setStartingOffsets(OffsetsInitializer.earliest())
                            .setProperties(consumProp)
                            .setValueOnlyDeserializer(new SimpleStringSchema())
                            .build();

            DataStreamSource<String> input =
                    env.fromSource(kfkSource, WatermarkStrategy.noWatermarks(), "raw_sa_event");

            SingleOutputStreamOperator<String> kafkaSource =
                    input.uid("raw_sa_event")
                            .filter(
                                    new FilterFunction<String>() {
                                        @Override
                                        public boolean filter(String record) {
                                            JSONObject jsonObject = JSONObject.parseObject(record);
                                            String message = jsonObject.getString("message");
                                            if (message != null) {
                                                return true;
                                            }
                                            return false;
                                        }
                                    })
                            .flatMap(
                                    new FlatMapFunction<String, String>() {
                                        @Override
                                        public void flatMap(String record, Collector<String> out) {
                                            try {

                                                /**
                                                 * post数据格式： log_format sa_extractor
                                                 * '"$proxy_add_x_forwarded_for" ++_ "$msec" ++_
                                                 * "$request_method" ++_ "$arg_gzip" ++_ "$arg_data"
                                                 * ++_ "$arg_data_list" ++_ "$request_body" ++_
                                                 * "$http_user_agent" ++_ "$arg_project" ++_
                                                 * "$http_cookie" ++_ "$arg_token" ++_ "$arg_ext"'
                                                 *
                                                 * <p>time时间校正： if (time is null, recv_time, if
                                                 * (_flush_time is null, time, if ( (recv_time -
                                                 * _flush_time) < 0 or (recv_time - _flush_time) >
                                                 * 60000, time + recv_time - _flush_time, time ) ) )
                                                 *
                                                 * <p>过滤条件： 1.event字段为null 2.nginx_time - time > 10d
                                                 * or time - nginx_time > 1h
                                                 */
                                                long jsonOutputTimestamp =
                                                        System.currentTimeMillis();
                                                JSONObject rawEventJson =
                                                        JSONObject.parseObject(record);
                                                String access_timestamp =
                                                        rawEventJson.getString("@timestamp");
                                                String message_urlEncode =
                                                        rawEventJson.getString("message");
                                                String[] data_format_arr =
                                                        message_urlEncode.split(" \\+\\+_ ");
                                                String recv_time =
                                                        data_format_arr[1].replace(
                                                                "\"", ""); // msec, nginx_time
                                                String request_method =
                                                        data_format_arr[2].replace(
                                                                "\"", ""); // request_method
                                                String project =
                                                        data_format_arr[8].replace(
                                                                "\"", ""); // project
                                                String request_body = null;
                                                if (request_method.equals("POST")) {
                                                    request_body =
                                                            data_format_arr[6].replace(
                                                                    "\"", ""); // post request_body
                                                } else if (request_method.equals("GET")) {
                                                    request_body =
                                                            data_format_arr[4].replace(
                                                                    "\"", ""); // get request_body
                                                }
                                                if (request_body != null) {
                                                    // URL解码
                                                    String request_base64Encode =
                                                            URLDecoder.decode(request_body);
                                                    if (request_method.equals("POST")
                                                            && request_base64Encode.contains(
                                                                    "data_list=")) {
                                                        String body_split =
                                                                request_base64Encode
                                                                        .split("data_list=")[1];
                                                        if (body_split.contains("&gzip=1")) {
                                                            body_split =
                                                                    body_split.split("&gzip=1")[0];
                                                        }
                                                        // Base64解码
                                                        byte[] body_gzipEncode =
                                                                Base64.getDecoder()
                                                                        .decode(body_split);
                                                        // Gzip解码
                                                        String dody_parse =
                                                                uncompressToString(
                                                                        body_gzipEncode,
                                                                        GZIP_ENCODE_UTF_8);
                                                        Object jsonData = JSON.parse(dody_parse);
                                                        if (jsonData instanceof JSONArray) {
                                                            JSONArray jsonArray =
                                                                    (JSONArray) jsonData;
                                                            for (Iterator<Object> iterator =
                                                                            jsonArray.iterator();
                                                                    iterator.hasNext(); ) {
                                                                JSONObject jsonObject =
                                                                        (JSONObject)
                                                                                iterator.next();
                                                                // 时间校正
                                                                Long nginx_time =
                                                                        Long.valueOf(
                                                                                recv_time.replace(
                                                                                        ".",
                                                                                        "")); // 秒级
                                                                Long flush_time =
                                                                        jsonObject.getLong(
                                                                                "_flush_time");
                                                                Long time =
                                                                        jsonObject.getLong("time");
                                                                Long timeAdjust =
                                                                        timeAdjust(
                                                                                nginx_time,
                                                                                flush_time,
                                                                                time);

                                                                // 数据过滤
                                                                if (jsonObject.getString("event")
                                                                                != null
                                                                        && (nginx_time - timeAdjust)
                                                                                <= 864000000
                                                                        && (timeAdjust - nginx_time)
                                                                                <= 3600000) {
                                                                    jsonObject.put(
                                                                            "time", timeAdjust);
                                                                    jsonObject.put(
                                                                            "access_timestamp",
                                                                            access_timestamp);
                                                                    jsonObject.put(
                                                                            "recv_time", recv_time);
                                                                    jsonObject.put(
                                                                            "access_type",
                                                                            request_method);
                                                                    jsonObject.put(
                                                                            "project", project);
                                                                    jsonObject.put(
                                                                            "json_output_timestamp",
                                                                            jsonOutputTimestamp);
                                                                    out.collect(
                                                                            jsonObject.toString(
                                                                                    SerializerFeature
                                                                                            .WriteMapNullValue));
                                                                } else {
                                                                    LOG.warn(
                                                                            "Filter data: "
                                                                                    + record);
                                                                }
                                                            }
                                                        } else if (jsonData instanceof JSONObject) {
                                                            JSONObject jsonObject =
                                                                    (JSONObject) jsonData;
                                                            // 时间校正
                                                            Long nginx_time =
                                                                    Long.valueOf(
                                                                            recv_time.replace(
                                                                                    ".", ""));
                                                            Long flush_time =
                                                                    jsonObject.getLong(
                                                                            "_flush_time");
                                                            Long time = jsonObject.getLong("time");
                                                            Long timeAdjust =
                                                                    timeAdjust(
                                                                            nginx_time,
                                                                            flush_time,
                                                                            time);

                                                            // 数据过滤
                                                            if (jsonObject.getString("event")
                                                                            != null
                                                                    && (nginx_time - timeAdjust)
                                                                            <= 864000000
                                                                    && (timeAdjust - nginx_time)
                                                                            <= 3600000) {
                                                                jsonObject.put("time", timeAdjust);
                                                                jsonObject.put(
                                                                        "access_timestamp",
                                                                        access_timestamp);
                                                                jsonObject.put(
                                                                        "recv_time", recv_time);
                                                                jsonObject.put(
                                                                        "access_type",
                                                                        request_method);
                                                                jsonObject.put("project", project);
                                                                jsonObject.put(
                                                                        "json_output_timestamp",
                                                                        jsonOutputTimestamp);
                                                                out.collect(
                                                                        jsonObject.toString(
                                                                                SerializerFeature
                                                                                        .WriteMapNullValue));
                                                            } else {
                                                                LOG.warn("Filter data: " + record);
                                                            }
                                                        } else {
                                                            LOG.warn(
                                                                    "No json instance data: "
                                                                            + jsonData);
                                                        }
                                                    } else if (request_method.equals("GET")) {
                                                        // Base64解码
                                                        byte[] body_parse =
                                                                Base64.getDecoder()
                                                                        .decode(
                                                                                request_base64Encode);
                                                        String body_str = new String(body_parse);
                                                        JSONObject jsonObject =
                                                                JSONObject.parseObject(body_str);
                                                        // 时间校正
                                                        Long nginx_time =
                                                                Long.valueOf(
                                                                        recv_time.replace(".", ""));
                                                        Long flush_time =
                                                                jsonObject.getLong("_flush_time");
                                                        Long time = jsonObject.getLong("time");
                                                        Long timeAdjust =
                                                                timeAdjust(
                                                                        nginx_time,
                                                                        flush_time,
                                                                        time);

                                                        // 数据过滤
                                                        if (jsonObject.getString("event") != null
                                                                && (nginx_time - timeAdjust)
                                                                        <= 864000000
                                                                && (timeAdjust - nginx_time)
                                                                        <= 3600000) {
                                                            jsonObject.put("time", timeAdjust);
                                                            jsonObject.put(
                                                                    "access_timestamp",
                                                                    access_timestamp);
                                                            jsonObject.put("recv_time", recv_time);
                                                            jsonObject.put(
                                                                    "access_type", request_method);
                                                            jsonObject.put("project", project);
                                                            jsonObject.put(
                                                                    "json_output_timestamp",
                                                                    jsonOutputTimestamp);
                                                            out.collect(
                                                                    jsonObject.toString(
                                                                            SerializerFeature
                                                                                    .WriteMapNullValue));
                                                        } else {
                                                            LOG.warn("Filter data: " + record);
                                                        }
                                                    } else {
                                                        LOG.warn(
                                                                "No GET or POST data: "
                                                                        + request_base64Encode);
                                                    }
                                                }
                                            } catch (Exception e) {
                                                LOG.warn("Error raw event data: " + record);
                                            }
                                        }
                                    });

            /** sink to Kafka */
            String sinkkfkTopic = paramReader.getKfkTopic();
            String bootStrapServer = paramReader.getKfkBootStrapServer();

            KafkaSink<String> kafkaSink =
                    KafkaSink.<String>builder()
                            .setBootstrapServers(bootStrapServer)
                            .setRecordSerializer(getSchemaBuild(sinkkfkTopic))
                            .build();

            kafkaSource.sinkTo(kafkaSink);

            env.execute(RawSaEventAccessLogParseFilter.class.toString());

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static KafkaRecordSerializationSchema<String> getSchemaBuild(String sinkkfkTopic) {
        return KafkaRecordSerializationSchema.builder()
                .setTopic(sinkkfkTopic)
                .setKeySerializationSchema(
                        new SerializationSchema<String>() {
                            @Override
                            public byte[] serialize(String record) {
                                JSONObject jsonObject = JSONObject.parseObject(record);
                                String keyby = jsonObject.getString(KEYBY);
                                return keyby.getBytes();
                            }
                        })
                .setValueSerializationSchema(new SimpleStringSchema())
                .setPartitioner(
                        new FlinkKafkaPartitioner<String>() {
                            @Override
                            public int partition(
                                    String record,
                                    byte[] key,
                                    byte[] value,
                                    String topic,
                                    int[] partitions) {
                                return Math.abs(new String(key).hashCode() % partitions.length);
                            }
                        })
                .build();
    }

    public static String uncompressToString(byte[] bytes, String encoding) {
        if (bytes == null || bytes.length == 0) {
            return null;
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ByteArrayInputStream in = new ByteArrayInputStream(bytes);
        try {
            GZIPInputStream ungzip = new GZIPInputStream(in);
            byte[] buffer = new byte[bytes.length];
            int n;
            while ((n = ungzip.read(buffer)) >= 0) {
                out.write(buffer, 0, n);
            }
            return out.toString(encoding);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * time时间校正： if (time is null, recv_time, if (_flush_time is null, time, if ( (recv_time -
     * _flush_time) < 0 or (recv_time - _flush_time) > 60000, time + recv_time - _flush_time, time )
     * ) )
     */
    public static Long timeAdjust(Long nginx_time, Long flush_time, Long time) {
        if (time == null) {
            return nginx_time;
        } else if (flush_time == null) {
            return time;
        } else if (nginx_time - flush_time < 0 || nginx_time - flush_time > 60000) {
            return time + nginx_time - flush_time;
        }
        return time;
    }
}
